package info.wangyuan.agent.service.rag;


import cn.hutool.core.codec.Base64;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import info.wangyuan.agent.common.Constant.PromptTemplate;
import info.wangyuan.agent.entity.dto.RagDocSyncChunkMsgDTO;
import info.wangyuan.agent.entity.po.FileInfo;
import info.wangyuan.agent.exception.BusinessException;
import info.wangyuan.agent.mapper.DocumentUnitMapper;
import info.wangyuan.agent.mapper.FileInfoMapper;
import info.wangyuan.agent.utils.PdfToBase64Converter;
import info.wangyuan.agent.utils.TikaFileTypeDetector;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * PDF文档切分策略实现
 *
 * @author Albert
 * @since 2025-08-27 17:48:29
 */
@Slf4j
@Service(value = "ragDocSyncOcr-PDF")
public abstract class PDFRagDocSyncChunkStrategyImpl implements RagDocSyncChunkStrategy {

    private final FileInfoMapper fileInfoMapper;
    private final DocumentUnitMapper documentUnitMapper;

    public PDFRagDocSyncChunkStrategyImpl(FileInfoMapper fileInfoMapper, DocumentUnitMapper documentUnitMapper) {
        this.fileInfoMapper = fileInfoMapper;
        this.documentUnitMapper = documentUnitMapper;
    }

    @Override
    public void handle(RagDocSyncChunkMsgDTO ragDocSyncChunkMsgDTO, String strategy) {
        Integer fileId = ragDocSyncChunkMsgDTO.getFileId();
        FileInfo fileInfo = fileInfoMapper.selectById(fileId);
        if (fileInfo == null) {
            log.error("文件ID {} 不存在", fileId);
            return;
        }

        log.info("开始对PDF文件进行切分，文件Id: {}", fileId);
        String filePath = fileInfo.getPath();
        byte[] fileBytes;
        try {
            fileBytes = fileToBytes(filePath);
            final int pdfPageCount = PdfToBase64Converter.getPdfPageCount(fileBytes);
            ragDocSyncChunkMsgDTO.setPageSize(pdfPageCount);

            // 更新数据库中的总页数
            LambdaUpdateWrapper<FileInfo> wrapper = Wrappers.lambdaUpdate(FileInfo.class)
                    .eq(FileInfo::getId, fileId)
                    .set(FileInfo::getPageSize, pdfPageCount);
            fileInfoMapper.update(wrapper);
            log.info("更新文件 {} 页数，总页数： {}", fileId, pdfPageCount);

            // 文档切分，OCR处理
            Map<Integer, String> ocrData = processFile(fileBytes, pdfPageCount, ragDocSyncChunkMsgDTO);

            // 文档片段入库
            insertData(ragDocSyncChunkMsgDTO, ocrData, documentUnitMapper);

            log.info("PDF文件切分入库完成，文件Id: {}", fileId);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 处理PDF文件 - 按页处理逻辑（带消息参数）
     */
    public Map<Integer, String> processFile(byte[] fileBytes,
                                            int totalPages,
                                            RagDocSyncChunkMsgDTO ragDocSyncChunkMsgDTO) {

        final HashMap<Integer, String> ocrData = new HashMap<>();
        for (int pageIndex = 0; pageIndex < totalPages; pageIndex++) {
            try {
                // 单独处理每一页以减少内存使用
                String base64 = PdfToBase64Converter.processPdfPageToBase64(fileBytes, pageIndex, "jpg");

                final UserMessage userMessage = UserMessage.userMessage(
                        ImageContent.from(base64, TikaFileTypeDetector.detectFileType(Base64.decode(base64))),
                        TextContent.from(PromptTemplate.RAG_OCR_PROMPT));

                ChatModel ocrModel = ragDocSyncChunkMsgDTO.getChatModel();
                if (ocrModel == null) {
                    throw new BusinessException("创建OCR模型失败");
                }

                final ChatResponse chat = ocrModel.chat(userMessage);

                ocrData.put(pageIndex, processText(chat.aiMessage().text()));

                if ((pageIndex + 1) % 10 == 0) {
                    System.gc();
                }

                try {
                    Thread.sleep(100);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            } catch (Exception e) {
                log.error("在处理PDF第 {} 页时发生错误: {}", (pageIndex + 1), e.getMessage());
                // 继续处理下一页，不中断整个流程
            }
        }
        return ocrData;
    }

    private static final Pattern[] PATTERNS = {
            Pattern.compile("\\\\（"),
            Pattern.compile("\\\\）"),
            Pattern.compile("\n{3,}"),
            Pattern.compile("([^\n])\n([^\n])"),
            Pattern.compile("\\$\\s+"),
            Pattern.compile("\\s+\\$"),
            Pattern.compile("\\$\\$")
    };

    public String processText(String input) {
        String result = input;
        result = PATTERNS[0].matcher(result).replaceAll(Matcher.quoteReplacement("\\("));
        result = PATTERNS[1].matcher(result).replaceAll(Matcher.quoteReplacement("\\)"));
        result = PATTERNS[2].matcher(result).replaceAll("\n\n");
        result = PATTERNS[3].matcher(result).replaceAll("$1\n$2");
        result = PATTERNS[4].matcher(result).replaceAll(Matcher.quoteReplacement("$"));
        result = PATTERNS[5].matcher(result).replaceAll(Matcher.quoteReplacement("$"));
        result = PATTERNS[6].matcher(result).replaceAll(Matcher.quoteReplacement("$$"));
        return result.trim();
    }
}
